#include "industrial_robot_control/kinematics/robot_arm.h"
#include "yaml-cpp/yaml.h"
#include <iostream>
#include <fstream>
#include <stdexcept>
#include <cmath>
#include <chrono>

namespace industrial_robot_control {
namespace kinematics {

RobotArm::RobotArm() : dof_(6), executing_trajectory_(false) {
    // 初始化状态向量
    current_state_.position.resize(dof_, 0.0);
    current_state_.velocity.resize(dof_, 0.0);
    current_state_.torque.resize(dof_, 0.0);
    
    target_state_.position.resize(dof_, 0.0);
    target_state_.velocity.resize(dof_, 0.0);
    target_state_.torque.resize(dof_, 0.0);
    
    // 默认关节限制
    joint_min_.resize(dof_, -M_PI);
    joint_max_.resize(dof_, M_PI);
    max_velocity_.resize(dof_, 1.0); // 1 rad/s
    max_acceleration_.resize(dof_, 2.0); // 2 rad/s²
}

bool RobotArm::initialize(const std::string& config_path) {
    try {
        if (!loadParameters(config_path)) {
            throw std::runtime_error("无法加载机器人参数");
        }
        
        std::cout << "机器人运动学初始化成功，自由度: " << dof_ << std::endl;
        return true;
    } catch (const std::exception& e) {
        std::cerr << "机器人运动学初始化错误: " << e.what() << std::endl;
        return false;
    }
}

void RobotArm::update() {
    // 计算当前笛卡尔空间位姿
    current_pose_ = forwardKinematics(current_state_.position);
    
    // 如果正在执行轨迹，更新目标状态
    if (executing_trajectory_ && !trajectory_.empty()) {
        std::lock_guard<std::mutex> lock(trajectory_mutex_);
        
        // 获取当前轨迹点
        auto current_point = trajectory_.front();
        
        // 更新目标状态
        target_state_ = current_point.state;
        
        // 移除已执行的轨迹点
        trajectory_.pop_front();
        
        // 如果轨迹执行完毕
        if (trajectory_.empty()) {
            executing_trajectory_ = false;
        }
    }
}

void RobotArm::sendCommands() {
    // 简单的PD控制器计算关节力矩
    for (int i = 0; i < dof_; ++i) {
        // 位置误差
        double pos_error = target_state_.position[i] - current_state_.position[i];
        
        // 速度误差
        double vel_error = target_state_.velocity[i] - current_state_.velocity[i];
        
        // PD控制律
        current_state_.torque[i] = 50.0 * pos_error + 5.0 * vel_error;
        
        // 限制力矩输出
        if (current_state_.torque[i] > 10.0) current_state_.torque[i] = 10.0;
        if (current_state_.torque[i] < -10.0) current_state_.torque[i] = -10.0;
        
        // 更新关节位置和速度 (简单模型)
        current_state_.velocity[i] += current_state_.torque[i] * 0.01; // 简化的动力学模型
        current_state_.position[i] += current_state_.velocity[i] * 0.01;
        
        // 限制关节位置在安全范围内
        if (current_state_.position[i] < joint_min_[i]) {
            current_state_.position[i] = joint_min_[i];
            current_state_.velocity[i] = 0.0;
        }
        if (current_state_.position[i] > joint_max_[i]) {
            current_state_.position[i] = joint_max_[i];
            current_state_.velocity[i] = 0.0;
        }
        
        // 限制关节速度
        if (current_state_.velocity[i] > max_velocity_[i]) {
            current_state_.velocity[i] = max_velocity_[i];
        }
        if (current_state_.velocity[i] < -max_velocity_[i]) {
            current_state_.velocity[i] = -max_velocity_[i];
        }
    }
}

CartesianPose RobotArm::forwardKinematics(const std::vector<double>& joint_angles) {
    if (joint_angles.size() != dof_) {
        throw std::invalid_argument("关节角度数量与自由度不匹配");
    }
    
    Eigen::Matrix4d transform = computeForwardKinematics(joint_angles);
    
    CartesianPose pose;
    pose.position.x() = transform(0, 3);
    pose.position.y() = transform(1, 3);
    pose.position.z() = transform(2, 3);
    
    Eigen::Matrix3d rotation = transform.block<3, 3>(0, 0);
    pose.orientation = Eigen::Quaterniond(rotation);
    
    return pose;
}

bool RobotArm::inverseKinematics(const CartesianPose& target, std::vector<double>& joint_angles) {
    // 简化的逆运动学实现，实际应用中需要更复杂的算法
    joint_angles.resize(dof_, 0.0);
    
    // 这里使用简单的数值迭代法求解逆解
    const int max_iterations = 1000;
    const double tolerance = 0.001; // 1mm位置误差
    
    // 从当前关节角度开始迭代
    joint_angles = current_state_.position;
    
    for (int iter = 0; iter < max_iterations; ++iter) {
        // 计算当前位姿
        CartesianPose current = forwardKinematics(joint_angles);
        
        // 计算位置误差
        Eigen::Vector3d pos_error = target.position - current.position;
        
        // 如果误差足够小，收敛成功
        if (pos_error.norm() < tolerance) {
            return true;
        }
        
        // 简单的雅可比矩阵伪逆方法更新关节角度
        // 实际应用中应使用更精确的雅可比矩阵计算
        for (int i = 0; i < dof_; ++i) {
            joint_angles[i] += 0.1 * pos_error[i % 3];
        }
        
        // 限制关节角度在范围内
        for (int i = 0; i < dof_; ++i) {
            if (joint_angles[i] < joint_min_[i]) joint_angles[i] = joint_min_[i];
            if (joint_angles[i] > joint_max_[i]) joint_angles[i] = joint_max_[i];
        }
    }
    
    // 迭代次数耗尽仍未收敛
    return false;
}

bool RobotArm::planJointTrajectory(
    const JointState& start, 
    const JointState& target, 
    double duration, 
    double sample_time) {
    
    if (start.position.size() != dof_ || target.position.size() != dof_) {
        return false;
    }
    
    if (duration <= 0 || sample_time <= 0 || sample_time > duration) {
        return false;
    }
    
    std::lock_guard<std::mutex> lock(trajectory_mutex_);
    trajectory_.clear();
    
    int num_points = static_cast<int>(duration / sample_time) + 1;
    
    // 创建轨迹点
    for (int i = 0; i < num_points; ++i) {
        double t = static_cast<double>(i) / (num_points - 1); // 0到1之间的归一化时间
        double time = t * duration;
        
        TrajectoryPoint point;
        point.time = time;
        point.state.position.resize(dof_);
        point.state.velocity.resize(dof_);
        point.state.torque.resize(dof_, 0.0);
        
        // 使用三次多项式插值
        for (int j = 0; j < dof_; ++j) {
            double p0 = start.position[j];
            double p1 = target.position[j];
            double v0 = start.velocity[j];
            double v1 = target.velocity[j];
            
            // 三次多项式系数
            double a0 = p0;
            double a1 = v0;
            double a2 = 3*(p1 - p0)/duration - 2*v0 - v1;
            double a3 = 2*(p0 - p1)/duration + v0 + v1;
            
            // 计算位置和速度
            point.state.position[j] = a0 + a1*t + a2*t*t + a3*t*t*t;
            point.state.velocity[j] = a1 + 2*a2*t + 3*a3*t*t;
            
            // 限制速度
            if (point.state.velocity[j] > max_velocity_[j]) {
                point.state.velocity[j] = max_velocity_[j];
            }
            if (point.state.velocity[j] < -max_velocity_[j]) {
                point.state.velocity[j] = -max_velocity_[j];
            }
        }
        
        trajectory_.push_back(point);
    }
    
    return true;
}

bool RobotArm::planCartesianTrajectory(
    const CartesianPose& start, 
    const CartesianPose& target, 
    double duration, 
    double sample_time) {
    
    // 将笛卡尔空间轨迹转换为关节空间轨迹
    JointState start_joint, target_joint;
    
    // 计算起始点和目标点的逆解
    if (!inverseKinematics(start, start_joint.position)) {
        return false;
    }
    if (!inverseKinematics(target, target_joint.position)) {
        return false;
    }
    
    // 速度初始化为0
    start_joint.velocity.assign(dof_, 0.0);
    target_joint.velocity.assign(dof_, 0.0);
    
    // 使用关节空间轨迹规划方法
    return planJointTrajectory(start_joint, target_joint, duration, sample_time);
}

void RobotArm::executeTrajectory() {
    if (!trajectory_.empty()) {
        executing_trajectory_ = true;
    }
}

void RobotArm::stopMotion() {
    std::lock_guard<std::mutex> lock(trajectory_mutex_);
    executing_trajectory_ = false;
    trajectory_.clear();
    
    // 平滑停止
    for (int i = 0; i < dof_; ++i) {
        target_state_.velocity[i] = 0.0;
    }
}

void RobotArm::emergencyStop() {
    std::lock_guard<std::mutex> lock(trajectory_mutex_);
    executing_trajectory_ = false;
    trajectory_.clear();
    
    // 立即停止所有运动
    for (int i = 0; i < dof_; ++i) {
        target_state_.position[i] = current_state_.position[i];
        target_state_.velocity[i] = 0.0;
        current_state_.velocity[i] = 0.0;
    }
    
    std::cout << "执行紧急停止" << std::endl;
}

void RobotArm::setJointLimits(const std::vector<double>& min, const std::vector<double>& max) {
    if (min.size() != dof_ || max.size() != dof_) {
        throw std::invalid_argument("关节限制数量与自由度不匹配");
    }
    
    joint_min_ = min;
    joint_max_ = max;
}

bool RobotArm::loadParameters(const std::string& config_path) {
    try {
        YAML::Node config = YAML::LoadFile(config_path);
        
        // 加载自由度
        if (config["dof"]) {
            dof_ = config["dof"].as<int>();
        }
        
        // 调整状态向量大小
        current_state_.position.resize(dof_);
        current_state_.velocity.resize(dof_);
        current_state_.torque.resize(dof_);
        
        target_state_.position.resize(dof_);
        target_state_.velocity.resize(dof_);
        target_state_.torque.resize(dof_);
        
        // 加载DH参数
        if (config["dh_parameters"]) {
            dh_params_ = config["dh_parameters"].as<std::vector<double>>();
        }
        
        // 加载关节限制
        if (config["joint_limits"]) {
            if (config["joint_limits"]["min"]) {
                joint_min_ = config["joint_limits"]["min"].as<std::vector<double>>();
            }
            if (config["joint_limits"]["max"]) {
                joint_max_ = config["joint_limits"]["max"].as<std::vector<double>>();
            }
            if (config["joint_limits"]["max_velocity"]) {
                max_velocity_ = config["joint_limits"]["max_velocity"].as<std::vector<double>>();
            }
            if (config["joint_limits"]["max_acceleration"]) {
                max_acceleration_ = config["joint_limits"]["max_acceleration"].as<std::vector<double>>();
            }
        }
        
        // 确保所有限制向量大小正确
        if (joint_min_.size() != dof_) joint_min_.resize(dof_, -M_PI);
        if (joint_max_.size() != dof_) joint_max_.resize(dof_, M_PI);
        if (max_velocity_.size() != dof_) max_velocity_.resize(dof_, 1.0);
        if (max_acceleration_.size() != dof_) max_acceleration_.resize(dof_, 2.0);
        
        return true;
    } catch (const std::exception& e) {
        std::cerr << "加载机器人参数错误: " << e.what() << std::endl;
        return false;
    }
}

Eigen::Matrix4d RobotArm::computeForwardKinematics(const std::vector<double>& joint_angles) {
    // 初始化变换矩阵为单位矩阵
    Eigen::Matrix4d transform = Eigen::Matrix4d::Identity();
    
    // 使用DH参数计算正运动学
    // 简化实现，实际应用中应根据机器人DH参数表计算
    for (int i = 0; i < dof_; ++i) {
        double theta = joint_angles[i];
        double d = 0.1 * (i + 1);  // 简化的连杆长度
        double a = 0.1 * (i + 1);  // 简化的连杆偏移
        double alpha = 0.0;
        
        // 如果有DH参数，使用配置的参数
        if (dh_params_.size() >= 4 * (i + 1)) {
            theta += dh_params_[4 * i];
            d = dh_params_[4 * i + 1];
            a = dh_params_[4 * i + 2];
            alpha = dh_params_[4 * i + 3];
        }
        
        // 计算单个关节的变换矩阵
        Eigen::Matrix4d joint_transform;
        joint_transform << cos(theta), -sin(theta)*cos(alpha),  sin(theta)*sin(alpha), a*cos(theta),
                          sin(theta),  cos(theta)*cos(alpha), -cos(theta)*sin(alpha), a*sin(theta),
                          0,           sin(alpha),             cos(alpha),            d,
                          0,           0,                      0,                     1;
        
        // 累积变换
        transform = transform * joint_transform;
    }
    
    return transform;
}

TrajectoryPoint RobotArm::interpolateTrajectory(const TrajectoryPoint& start, 
                                              const TrajectoryPoint& end, 
                                              double t) {
    TrajectoryPoint result;
    result.time = start.time + t * (end.time - start.time);
    result.state.position.resize(dof_);
    result.state.velocity.resize(dof_);
    result.state.torque.resize(dof_);
    
    // 线性插值
    for (int i = 0; i < dof_; ++i) {
        result.state.position[i] = start.state.position[i] + t * (end.state.position[i] - start.state.position[i]);
        result.state.velocity[i] = start.state.velocity[i] + t * (end.state.velocity[i] - start.state.velocity[i]);
        result.state.torque[i] = start.state.torque[i] + t * (end.state.torque[i] - start.state.torque[i]);
    }
    
    return result;
}

} // namespace kinematics
} // namespace industrial_robot_control
